import torch

from ignite.metrics import Metric

from ignite.exceptions import NotComputableError
from ignite.metrics.metric import sync_all_reduce, reinit__is_reduced

class ENL(Metric):


    def __init__(self, output_transform=lambda x: x, device="cpu", thresh=[1, 2, 5]):
        self._sum = None
        self._num = None
        super(ENL, self).__init__(output_transform=output_transform, device=device)


    @reinit__is_reduced
    def reset(self):
        self._sum = 0
        self._num = 0
        super().reset()


    @reinit__is_reduced
    def update(self, output):
        despeckle, _ = output

        b = despeckle.size(0)
        self._num += b

        roi = despeckle.view(b, -1)
        mu = torch.mean(roi, dim=1)
        std = torch.std(roi, dim=1, correction=0)
        enl = (mu * mu / std / std).sum()
        self._sum += enl


    @sync_all_reduce("_num", "_sum:SUM")
    def compute(self):
        if self._num == 0:
            raise NotComputableError("ENL has 0 sample!")

        enl = self._sum / self._num

        return enl


if __name__ == "__main__":
    m = ENL()

    img = torch.randint(255, (5, 1, 256, 256)).float()

    for i in range(img.size(0)):
        m.update((img[i], None))
    
    res = m.compute()
    print(res)


